import pandas as pd
import os
import shutil
import random
from tqdm import tqdm

# 读取 CSV 文件
# file_path = 'training_anno.csv'  # CSV 文件路径
# output_dir = './10w'  # 新的数据集存放路径
# train_dir = './10w/train'  # 训练集存放路径
# val_dir = './10w/val'  # 验证集存放路径
# image_dir = './'  # 原始图片存放路径

file_path = 'training_anno.csv'  # CSV 文件路径
output_dir = './trainval'  # 新的数据集存放路径
train_dir = './trainval/train'  # 训练集存放路径
val_dir = './trainval/val'  # 验证集存放路径
image_dir = './'  # 原始图片存放路径

# 将 CSV 文件读取为 DataFrame
df = pd.read_csv(file_path)

# 确认数据集有图片路径列
image_column = 'Path'  # 修改为实际的图片路径列名

# 抽取四分之一的数据
# sampled_df = df.sample(frac=0.01, random_state=42)
sampled_df = df

# 按照9比1的比例划分为训练集和验证集
train_df = sampled_df.sample(frac=0.9, random_state=42)
val_df = sampled_df.drop(train_df.index)

# 创建新的数据集文件夹
if not os.path.exists(train_dir):
    os.makedirs(train_dir)

if not os.path.exists(val_dir):
    os.makedirs(val_dir)

# 保存训练集和验证集的标注文件
train_df.to_csv(os.path.join(train_dir, 'train.csv'), index=False)
val_df.to_csv(os.path.join(val_dir, 'val.csv'), index=False)

# 复制训练集图片到新的文件夹，并保持目录结构，使用 tqdm 显示进度条
for idx, row in tqdm(train_df.iterrows(), total=train_df.shape[0], desc="Copying Training Images"):
    # 原始图片完整路径
    src_image_path = os.path.join(image_dir, row[image_column])

    if os.path.exists(src_image_path):
        # 目标图片路径，包括文件夹结构
        relative_path = os.path.relpath(row[image_column], image_dir)  # 计算相对路径
        dest_image_path = os.path.join(train_dir, relative_path)

        # 创建目标文件夹（如果不存在）
        os.makedirs(os.path.dirname(dest_image_path), exist_ok=True)

        # 复制文件
        shutil.copy(src_image_path, dest_image_path)

# 复制验证集图片到新的文件夹，并保持目录结构，使用 tqdm 显示进度条
for idx, row in tqdm(val_df.iterrows(), total=val_df.shape[0], desc="Copying Validation Images"):
    # 原始图片完整路径
    src_image_path = os.path.join(image_dir, row[image_column])

    if os.path.exists(src_image_path):
        # 目标图片路径，包括文件夹结构
        relative_path = os.path.relpath(row[image_column], image_dir)  # 计算相对路径
        dest_image_path = os.path.join(val_dir, relative_path)

        # 创建目标文件夹（如果不存在）
        os.makedirs(os.path.dirname(dest_image_path), exist_ok=True)

        # 复制文件
        shutil.copy(src_image_path, dest_image_path)

print("数据集按照9比1划分并保存完毕！")
